import sys
import random
import argparse
import json
import pathlib
import numpy as np
from copy import deepcopy

from PyQt5 import QtGui
from PyQt5.QtWidgets import (
    QApplication, QWidget, QPushButton, QLabel, QTextEdit, QLineEdit,
    QVBoxLayout, QHBoxLayout, QFormLayout, QMessageBox, QFileDialog
)
from PyQt5.QtCore import Qt, QObject, pyqtSignal

# 假设 src 目录已正确导入
from scripts.reconst import compare_graphs
from src.core.PartType import PartType, random_part_type, random_pipe_type
from src.graph_creator.core.PipelineGraph import PipelineGraph
from src.classifier.trainer.TrainingModel import create_training_model, TrainingModel
from src.graph_creator.GraphCreator import GraphCreator
from src.pipeline_constructor.PipelineConstructor import PipelineConstructor
import open3d as o3d

from src.pipeline_constructor.core.PipelineModel import PipelineModel


class Logger(QObject):
    message_written = pyqtSignal(str)

    def write(self, msg):
        self.message_written.emit(msg)

    def flush(self):
        pass


class PipelineReconstructionGUI(QWidget):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("管道结构重建工具")
        self.resize(800, 600)

        # 初始化变量
        self.model = None
        self.reconstructed_model = None
        self.save_path = './output'
        self.resource_dir = pathlib.Path(__file__).parent
        self.init_ui()

    def init_ui(self):
        layout = QVBoxLayout()

        # 参数输入表单
        form_layout = QFormLayout()
        self.nb_parts_input = QLineEdit("10")
        self.sampling_input = QLineEdit("100")
        self.accuracy_input = QLineEdit("0.95")
        self.center_error_input = QLineEdit("0.05")
        self.direction_error_input = QLineEdit("0.05")
        self.save_name_input = QLineEdit("ReconstructedGraph")

        form_layout.addRow("部件数量:", self.nb_parts_input)
        form_layout.addRow("采样点数:", self.sampling_input)
        form_layout.addRow("分类准确率 (0~1):", self.accuracy_input)
        form_layout.addRow("中心误差比例 (0~1):", self.center_error_input)
        form_layout.addRow("方向误差比例 (0~1):", self.direction_error_input)
        form_layout.addRow("保存文件名:", self.save_name_input)

        layout.addLayout(form_layout)

        # 操作按钮
        button_layout = QHBoxLayout()
        self.select_file_btn = QPushButton("选择JSON文件")
        self.generate_btn = QPushButton("生成随机管道")
        self.reconstruct_btn = QPushButton("重建管道")
        self.visualize_btn = QPushButton("可视化重建结果")
        self.save_btn = QPushButton("保存重建模型")
        self.compare_btn = QPushButton("对比可视化")

        self.select_file_btn.clicked.connect(self.select_ply_file)
        self.generate_btn.clicked.connect(self.generate_pipeline)
        self.reconstruct_btn.clicked.connect(self.reconstruct_pipeline)
        self.visualize_btn.clicked.connect(self.visualize_reconstructed)
        self.save_btn.clicked.connect(self.save_reconstructed)
        self.compare_btn.clicked.connect(self.compare_models)

        button_layout.addWidget(self.select_file_btn)
        button_layout.addWidget(self.generate_btn)
        button_layout.addWidget(self.reconstruct_btn)
        button_layout.addWidget(self.visualize_btn)
        button_layout.addWidget(self.save_btn)
        button_layout.addWidget(self.compare_btn)

        layout.addLayout(button_layout)

        # 日志输出区域
        self.log_text = QTextEdit()
        self.log_text.setReadOnly(True)
        layout.addWidget(QLabel("操作日志:"))
        layout.addWidget(self.log_text)

        self.setLayout(layout)

        # 重定向 stdout 到日志框
        sys.stdout = Logger(message_written=self.append_log)


    def select_ply_file(self):
        options = QFileDialog.Options()
        file_path, _ = QFileDialog.getOpenFileName(
            self,
            "选择 JSON 文件",
            "",
            "JSON Files (*.json);;All Files (*)",
            options=options
        )
        if file_path:
            print(f"✅ 已选择文件: {file_path}")
            self.load_and_use_ply(file_path)

    def load_and_use_ply(self, file_path):
        try:
            print("🔍 正在显示所选点云...【JSON格式保存】")
            # 封装为 TrainingModel
            self.model = TrainingModel(json_file=file_path)
            self.model.visualize()
            print("✅ 点云已成功加载并封装为 TrainingModel！")
        except Exception as e:
            QMessageBox.critical(self, "错误", f"加载失败: {str(e)}")
            return  # ✅ 提前退出，防止后续操作 None 类型


    def append_log(self, msg):
        self.log_text.moveCursor(QtGui.QTextCursor.End)
        self.log_text.insertPlainText(msg)
        self.log_text.ensureCursorVisible()

    def generate_pipeline(self):
        try:
            nb_parts = int(self.nb_parts_input.text())
            sampling = int(self.sampling_input.text())
            print(f"生成 {nb_parts} 个部件，每个部件 {sampling} 个采样点...")
            self.model = create_training_model(nb_parts=nb_parts, nb_points_per_mesh=sampling)
            print("✅ 原始管道生成完成！")
            self.model.visualize()
        except Exception as e:
            QMessageBox.critical(self, "错误", f"生成失败: {str(e)}")

    def reconstruct_pipeline(self):
        if not self.model:
            QMessageBox.warning(self, "警告", "请先生成管道！")
            return

        try:
            args_accuracy = float(self.accuracy_input.text())
            args_center_error = float(self.center_error_input.text())
            args_direction_error = float(self.direction_error_input.text())

            classification = deepcopy(self.model.point_cloud_unlabelled)
            labels_truth_points = deepcopy(self.model.point_cloud_labelled.points)

            # for i in range(len(classification.points)):
            #     ground_truth = self.model.point_cloud_labelled.points[i]
            #     prob = random.randrange(0, 100)
            #     if prob >= args_accuracy * 100:
            #         error_part_type = ground_truth.part_type
            #         while error_part_type == ground_truth.part_type:
            #             error_part_type = random_pipe_type() if ground_truth.part_type.name.startswith(
            #                 'PIPE') else random_part_type()
            #         classification.points[i].part_type = error_part_type
            #     else:
            #         classification.points[i].part_type = ground_truth.part_type
            #
            #     center_error = random.uniform(-args_center_error, args_center_error)
            #     direction_error = random.uniform(-args_direction_error, args_direction_error)
            #     classification.points[i].center = ground_truth.center + ground_truth.center * center_error
            #     classification.points[i].direction = ground_truth.direction + ground_truth.direction * direction_error
            #     classification.points[i].radius = 1.
            for i in range(len(classification.points)):
                ground_truth = self.model.point_cloud_labelled.points[i]

                prob = random.randrange(0, 100)
                if prob >= args_accuracy * 100:
                    error_part_type = ground_truth.part_type
                    while error_part_type == ground_truth.part_type:
                        error_part_type = random_pipe_type() if ground_truth.part_type.name.startswith(
                            'PIPE') else random_part_type()
                    classification.points[i].part_type = error_part_type
                else:
                    classification.points[i].part_type = ground_truth.part_type

                # ✅ 使用 NumPy 处理 center 和 direction
                center_error = random.uniform(-args_center_error, args_center_error)
                classification.points[i].center = np.array(ground_truth.center) * (1 + center_error)

                direction_error = random.uniform(-args_direction_error, args_direction_error)
                classification.points[i].direction = np.array(ground_truth.direction) * (1 + direction_error)

                classification.points[i].radius = 1.

            print("🔍 正在构建图结构...")
            graph_creator = GraphCreator()
            pipeline_constructor = PipelineConstructor()
            reconstructed_graph = graph_creator.build_graph(classification)
            self.reconstructed_model = pipeline_constructor.construct_pipeline(reconstructed_graph)

            if self.model.model_graph:
                compare_graphs(self.model.model_graph, reconstructed_graph)

            print("✅ 重建完成！")
        except Exception as e:
            QMessageBox.critical(self, "错误", f"重建失败: {str(e)}")

    def visualize_reconstructed(self):
        if not self.reconstructed_model:
            QMessageBox.warning(self, "警告", "请先完成重建！")
            return
        try:
            self.reconstructed_model.visualize()
        except Exception as e:
            QMessageBox.critical(self, "错误", f"可视化失败: {str(e)}")

    def save_reconstructed(self):
        if not self.reconstructed_model:
            QMessageBox.warning(self, "警告", "请先完成重建！")
            return
        try:
            file_name = self.save_name_input.text()
            file_path, _ = QFileDialog.getSaveFileName(self, "保存重建模型", file_name, "PLY 文件 (*.ply);;OBJ 文件 (*.obj)")
            if file_path:
                ext = pathlib.Path(file_path).suffix or ".ply"
                self.reconstructed_model.save(file_path, extension=ext)
                print(f"💾 模型已保存至: {file_path}")
        except Exception as e:
            QMessageBox.critical(self, "错误", f"保存失败: {str(e)}")

    def compare_models(self):
        if not self.model or not self.reconstructed_model:
            QMessageBox.warning(self, "警告", "请先生成并重建管道！")
            return
        try:
            self._visualize_two_models(self.model, self.reconstructed_model)
            # self._visualize_two_models_in_one_window(self.model, self.reconstructed_model)
        except Exception as e:
            QMessageBox.critical(self, "错误", f"对比失败: {str(e)}")

    def _visualize_two_models(self, original_model, reconstructed_model):
        """
        同时可视化原始模型和重建模型，支持不同几何类型并保持视角同步。
        """

        # 获取几何体
        if isinstance(original_model, TrainingModel):
            geom1 = original_model.point_cloud_model  # 原始是点云
        elif hasattr(original_model, 'get_visualization_data'):
            geom1 = original_model.get_visualization_data()
        else:
            geom1 = original_model.mesh  # 或者尝试用 mesh

        if isinstance(reconstructed_model, PipelineModel):
            geom2 = reconstructed_model.mesh  # 重建的是网格
        elif hasattr(reconstructed_model, 'get_visualization_data'):
            geom2 = reconstructed_model.get_visualization_data()
        else:
            raise TypeError("无法识别的重建模型类型")

        # 边界框
        box1 = geom1.get_axis_aligned_bounding_box()
        box1.color = (1, 0, 0)
        box2 = geom2.get_axis_aligned_bounding_box()
        box2.color = (0, 1, 0)

        # 坐标系
        frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=1.0)

        # 创建两个可视化器
        vis1 = o3d.visualization.Visualizer()
        vis1.create_window(window_name='原始管道', width=600, height=600, left=0, top=0)
        vis1.add_geometry(geom1)
        vis1.add_geometry(box1)
        vis1.add_geometry(frame)

        vis2 = o3d.visualization.Visualizer()
        vis2.create_window(window_name='重建管道', width=600, height=600, left=610, top=0)
        vis2.add_geometry(geom2)
        vis2.add_geometry(box2)
        vis2.add_geometry(frame)

        # 同步相机参数
        def synchronize_views(vis1, vis2):
            view_ctl1 = vis1.get_view_control()
            view_ctl2 = vis2.get_view_control()

            # 复制相机参数
            cam_param = view_ctl1.convert_to_pinhole_camera_parameters()
            view_ctl2.convert_from_pinhole_camera_parameters(cam_param)

        # 更新函数（使用 poll_events 避免死循环）
        def update():
            while True:
                vis1.poll_events()
                vis2.poll_events()

                synchronize_views(vis1, vis2)

                vis1.update_renderer()
                vis2.update_renderer()

                # 如果任意一个窗口关闭，停止更新
                if not vis1.poll_events() or not vis2.poll_events():
                    break

        # 开始更新
        try:
            update()
        except Exception as e:
            print(f"可视化过程中出错: {e}")
        finally:
            vis1.destroy_window()
            vis2.destroy_window()

    def merge_bounding_boxes(self, box1, box2):
        min_bound = np.minimum(box1.get_min_bound(), box2.get_min_bound())
        max_bound = np.maximum(box1.get_max_bound(), box2.get_max_bound())
        merged_box = o3d.geometry.AxisAlignedBoundingBox(min_bound, max_bound)
        return merged_box

    def _visualize_two_models_in_one_window(self, original_model, reconstructed_model):
        """
        在一个窗口中同时显示原始和重建模型，支持点云和网格。
        """

        # 获取几何体
        if isinstance(original_model, TrainingModel):
            geom1 = original_model.point_cloud_model
        elif hasattr(original_model, 'visualization_data'):
            geom1 = original_model.visualization_data
        else:
            raise TypeError("无法识别的原始模型类型")

        if isinstance(reconstructed_model, PipelineModel):
            geom2 = reconstructed_model.mesh
        elif hasattr(reconstructed_model, 'visualization_data'):
            geom2 = reconstructed_model.visualization_data
        else:
            raise TypeError("无法识别的重建模型类型")

        # 对其中一个模型进行平移，使其与另一个并列显示
        trans_x_offset = 5.0  # 水平偏移量，单位：米

        def shift_geometry(geom, offset):
            if isinstance(geom, o3d.geometry.PointCloud):
                points = np.asarray(geom.points)
                points[:, 0] += offset
                geom.points = o3d.utility.Vector3dVector(points)
            elif isinstance(geom, o3d.geometry.TriangleMesh):
                vertices = np.asarray(geom.vertices)
                vertices[:, 0] += offset
                geom.vertices = o3d.utility.Vector3dVector(vertices)
            return geom

        geom1 = shift_geometry(geom1, -trans_x_offset)
        geom2 = shift_geometry(geom2, trans_x_offset)

        # 设置颜色区分
        if isinstance(geom1, o3d.geometry.PointCloud):
            geom1.paint_uniform_color([0.1, 0.1, 0.7])  # 原始模型：蓝色

        if isinstance(geom2, o3d.geometry.PointCloud):
            geom2.paint_uniform_color([0.7, 0.1, 0.1])  # 重建模型：红色

        if isinstance(geom1, o3d.geometry.TriangleMesh):
            geom1.paint_uniform_color([0.3, 0.3, 1.0])  # 网格浅蓝

        if isinstance(geom2, o3d.geometry.TriangleMesh):
            geom2.paint_uniform_color([1.0, 0.3, 0.3])  # 网格浅红

        # 合并包围盒
        box1 = geom1.get_axis_aligned_bounding_box()
        box2 = geom2.get_axis_aligned_bounding_box()
        box_combined = self.merge_bounding_boxes(box1, box2)
        box_combined.color = (0, 0, 0)  # 黑色边框

        # 添加坐标系
        frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=1.0)

        # 组合所有几何体
        geometries = [geom1, geom2, frame, box_combined]

        # 可视化
        print("🔍 正在显示对比结果（关闭窗口继续操作）...")
        o3d.visualization.draw_geometries(
            geometries,
            window_name="原始 vs 重建管道",
            width=1200,
            height=800
        )


if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = PipelineReconstructionGUI()
    window.show()
    sys.exit(app.exec_())